import json
import rsa
import base64
from Crypto.Cipher import AES
from core.tools.custom_requests import CustomRequests

BLOCK_SIZE = AES.block_size
pad = lambda s, length: s + (BLOCK_SIZE - length % BLOCK_SIZE) * chr(BLOCK_SIZE - length % BLOCK_SIZE)
unpad = lambda s: s[:-ord(s[len(s) - 1:])]


class BaseAliPayClient(CustomRequests):
    PYTHON_VERSION_3 = True

    @property
    def get_charset(self):
        return getattr(self, "charset", "utf-8")

    @property
    def get_private_key(self):
        private_key = getattr(self, "private_key", None)
        if private_key:
            return private_key
        else:
            raise Exception("没有配置应用私钥")

    @property
    def get_public_key(self):
        private_key = getattr(self, "public_key", None)
        if private_key:
            return private_key
        else:
            raise Exception("没有配置应用公钥")

    @property
    def get_encrypt_key(self):
        private_key = getattr(self, "encrypt_key", None)
        if private_key:
            return private_key
        else:
            raise Exception("没有配置AES文件")

    @property
    def get_ali_public_key(self):
        private_key = getattr(self, "ali_public_key", None)
        if private_key:
            return private_key
        else:
            raise Exception("没有配置支付宝公钥")

    @classmethod
    def get_sign_content(cls, all_params):
        sign_content = ""
        for (k, v) in sorted(all_params.items()):
            value = v
            if not isinstance(value, str):
                value = json.dumps(value, ensure_ascii=False)
            sign_content += ("&" + k + "=" + value)
        sign_content = sign_content[1:]
        return sign_content

    @classmethod
    def add_start_end(cls, key, startMarker, endMarker):
        if key.find(startMarker) < 0:
            key = startMarker + key
        if key.find(endMarker) < 0:
            key = key + endMarker
        return key

    @classmethod
    def fill_private_key_marker(cls, private_key):
        return cls.add_start_end(private_key, "-----BEGIN RSA PRIVATE KEY-----\n", "\n-----END RSA PRIVATE KEY-----")

    @classmethod
    def fill_public_key_marker(cls, public_key):
        return cls.add_start_end(public_key, "-----BEGIN PUBLIC KEY-----\n", "\n-----END PUBLIC KEY-----")

    def sign_with_rsa(self, sign_content, charset):
        if self.PYTHON_VERSION_3:
            sign_content = sign_content.encode(charset)
        private_key = self.fill_private_key_marker(self.get_private_key)
        signature = rsa.sign(sign_content, rsa.PrivateKey.load_pkcs1(private_key, format='PEM'), 'SHA-1')
        sign = base64.b64encode(signature)
        if self.PYTHON_VERSION_3:
            sign = str(sign, encoding=charset)
        return sign

    # 使用应用私钥签名
    def sign_with_rsa2(self, sign_content):
        if self.PYTHON_VERSION_3:
            sign_content = sign_content.encode(self.get_charset)
        private_key = self.fill_private_key_marker(self.get_private_key)
        signature = rsa.sign(sign_content, rsa.PrivateKey.load_pkcs1(private_key, format='PEM'), 'SHA-256')
        sign = base64.b64encode(signature)
        if self.PYTHON_VERSION_3:
            sign = str(sign, encoding=self.get_charset)
        return sign

    # 使用支付宝公钥验签
    def verify_with_rsa(self, message, sign):
        public_key = self.fill_public_key_marker(self.get_ali_public_key)
        sign = base64.b64decode(sign)
        return bool(rsa.verify(message, sign, rsa.PublicKey.load_pkcs1_openssl_pem(public_key)))

    # 加密
    def encrypt_content(self, content, encrypt_type):
        if "AES" == encrypt_type.upper():
            return self.aes_encrypt_content(content)
        raise Exception("当前不支持该算法类型encrypt_type=" + encrypt_type)

    def aes_encrypt_content(self, content):
        if self.PYTHON_VERSION_3:
            iv = b'\0' * BLOCK_SIZE
            length = len(bytes(content, encoding=self.get_charset))
        else:
            length = len(bytes(content))
            iv = '\0' * BLOCK_SIZE
        padded_content = pad(content, length)
        cryptor = AES.new(base64.b64decode(self.get_encrypt_key), AES.MODE_CBC, iv)
        encrypted_content = cryptor.encrypt(padded_content)
        encrypted_content = base64.b64encode(encrypted_content)
        if self.PYTHON_VERSION_3:
            encrypted_content = str(encrypted_content, encoding=self.get_charset)
        return encrypted_content

    # 解密
    def decrypt_content(self, encrypted_content, encrypt_type):
        if "AES" == encrypt_type.upper():
            return self.aes_decrypt_content(encrypted_content)
        raise Exception("当前不支持该算法类型encrypt_type=" + encrypt_type)

    def aes_decrypt_content(self, encrypted_content):
        encrypted_content = base64.b64decode(encrypted_content)
        if self.PYTHON_VERSION_3:
            iv = b'\0' * BLOCK_SIZE
        else:
            iv = '\0' * BLOCK_SIZE
        cryptor = AES.new(base64.b64decode(self.get_encrypt_key), AES.MODE_CBC, iv)
        content = unpad(cryptor.decrypt(encrypted_content))
        if self.PYTHON_VERSION_3:
            content = content.decode(self.get_charset)
        return content
